{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0927e511-be9f-4fcf-84c5-f7b9c37f8179",
   "metadata": {},
   "source": [
    "# Census 2020 Visualization using Plotly-Dash + RAPIDS on Google Colab\n",
    "\n",
    "<img src=\"../assets/dashboard.png\" style=\"max-width:auto; height:500px;margin-left:30%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be4d30b4",
   "metadata": {},
   "source": [
    "1. [Plotly Dash](https://plotly.com/dash/) is a Python framework for building web applications. It is built on top of Flask, Plotly.js, React and React Js. Dash is ideal for building data visualization apps with highly custom user interfaces in pure Python. It is particularly suited for anyone who works with data in Python.\n",
    "2. [RAPIDS](https://rapids.ai/) is a suite of open-source libraries for end-to-end data science pipelines entirely in the GPU. It is designed to have a familiar look and feel to data scientists working in Python. RAPIDS is based on the Apache Arrow columnar memory format, and cuDF is a GPU DataFrame library for loading, joining, aggregating, filtering, and otherwise manipulating data.\n",
    "3. [Google Colab](https://colab.research.google.com/) is a Jupyter notebook like environment that requires no setup to use and runs entirely in the cloud. It is a free service and you can use it to run Python and R code. It also provides free GPU to run your code.\n",
    "\n",
    "This notebook demonstrates how to use Plotly Dash with RAPIDS in Google Colab to visualize the Census 2020 data. The Census 2020 data is available in the form of parquet files. The CSV files are loaded into GPU memory using cuDF. The data is then visualized using Plotly Dash. The visualization is interactive and the user can select the state and the county to view the population data, migration data from 2010 to 2020 and the population density. The data is also visualized on a map using Plotly Scattermapbox."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "535ccd22",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "578b4edb",
   "metadata": {},
   "source": [
    "### Install RAPIDS cuDF and Plotly Dash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "zAUP3v24hdU2",
   "metadata": {
    "id": "zAUP3v24hdU2"
   },
   "outputs": [],
   "source": [
    "!pip install cudf-cu11 --extra-index-url=https://pypi.nvidia.com\n",
    "!pip install dash dash_bootstrap_components jupyter-dash dash-html-components dash-core-components dash-daq dash-bootstrap-components \"datashader>=0.15\" pyproj bokeh"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdee70ca",
   "metadata": {},
   "source": [
    "### Download the Census 2020 data\n",
    "\n",
    "The Total Population dataset includes Census 2020 total population data with decennial migration from Census 2010 at a block level.  Please note that due to GPU memory limitations (the T4 GPU available via Google Colab has 16GB of memory), the Total Population dataset, which has more than 381M rows, may sometimes result in Out of Memory Errors. Using individual state data, which has fewer rows, is recommended.\n",
    "\n",
    "For more information on how the Census 2020 and 2010 Migration data was prepared to show individual points, please refer to the [Data preparation notebooks](https://github.com/rapidsai/plotly-dash-rapids-census-demo/tree/main/data_prep_total_population)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3wPfCzvBDflt",
   "metadata": {
    "id": "3wPfCzvBDflt"
   },
   "outputs": [],
   "source": [
    "# Download state-wise datasets\n",
    "import os\n",
    "\n",
    "if not os.path.exists(\"state-wise-population\"):\n",
    "    !wget https://data.rapids.ai/viz-data/state-wise-population.tar.xz\n",
    "    !tar -xJf state-wise-population.tar.xz\n",
    "    !rm state-wise-population.tar.xz\n",
    "\n",
    "# Optional: Download all US states combined dataset, uncomment the following lines to download the dataset\n",
    "# if not os.path.exists(\"total_population_dataset.parquet\"):\n",
    "#     !wget https://data.rapids.ai/viz-data/total_population_dataset.parquet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6382f95c",
   "metadata": {},
   "source": [
    "### Download assets and utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceb8334a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download CSS and logo files from the original repo to use in the plotly dashboard\n",
    "! mkdir -p assets && for file in s1.css dash-logo.png rapids-logo.png; do curl -o assets/$file https://raw.githubusercontent.com/rapidsai/plotly-dash-rapids-census-demo/main/plotly_demo/assets/$file; done\n",
    "\n",
    "# Download utils.py file from the original repo to use the utility functions in the notebook\n",
    "! mkdir -p utils && for file in __init__.py utils.py; do curl -o utils/$file https://raw.githubusercontent.com/rapidsai/plotly-dash-rapids-census-demo/main/plotly_demo/utils/$file; done"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8658b7d",
   "metadata": {},
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "115d0fdb-428d-49f9-858f-424289043d07",
   "metadata": {
    "id": "115d0fdb-428d-49f9-858f-424289043d07"
   },
   "outputs": [],
   "source": [
    "from bokeh import palettes\n",
    "from dash import Dash, ctx, dcc, html\n",
    "from dash.dependencies import Input, Output, State\n",
    "from dash.exceptions import PreventUpdate\n",
    "from distributed import Client\n",
    "from pyproj import Transformer\n",
    "import cupy as cp\n",
    "import dash_bootstrap_components as dbc\n",
    "import dash_daq as daq\n",
    "import datashader as ds\n",
    "import datashader.transfer_functions as tf\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pickle\n",
    "import requests\n",
    "import time\n",
    "\n",
    "from jupyter_dash import JupyterDash # For running the app in a notebook environment    \n",
    "import cudf # RAPIDS cuDF is an implementation of Pandas-like Dataframe on GPU\n",
    "\n",
    "from utils import * # Import utility functions from the utils.py file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4179466d",
   "metadata": {},
   "source": [
    "### Set Dashboard Attributes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5be83f53-7b40-4256-a7e3-dd436bb8a697",
   "metadata": {
    "id": "5be83f53-7b40-4256-a7e3-dd436bb8a697"
   },
   "outputs": [],
   "source": [
    "DATA_PATH = \".\"\n",
    "DATA_PATH_STATE = f\"{DATA_PATH}/state-wise-population\"\n",
    "\n",
    "# Get list of state names from file names\n",
    "state_files = os.listdir(DATA_PATH_STATE)\n",
    "state_names = [os.path.splitext(f)[0] for f in state_files]\n",
    "\n",
    "DEFAULT_STATE = state_names[0]\n",
    "\n",
    "# append USA (combined) to state_names if it exists\n",
    "if os.path.exists(\"total_population_dataset.parquet\"):\n",
    "    state_names.append(\"USA\")\n",
    "\n",
    "(\n",
    "    data_center_3857,\n",
    "    data_3857,\n",
    "    data_4326,\n",
    "    data_center_4326,\n",
    "    selected_map_backup,\n",
    "    selected_race_backup,\n",
    "    selected_county_top_backup,\n",
    "    selected_county_bt_backup,\n",
    "    view_name_backup,\n",
    "    c_df,\n",
    "    gpu_enabled_backup,\n",
    "    dragmode_backup,\n",
    "    currently_loaded_state,\n",
    ") = ([], [], [], [], None, None, None, None, None, None, None, \"pan\", None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9da7be00",
   "metadata": {},
   "source": [
    "### Initialize the Dash app via JupyterDash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "IWV-3SB-C3pw",
   "metadata": {
    "id": "IWV-3SB-C3pw"
   },
   "outputs": [],
   "source": [
    "app = JupyterDash(__name__)\n",
    "# Create server variable with Flask server object for use with gunicorn\n",
    "server = app.server"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "007eb387",
   "metadata": {},
   "source": [
    "### Dashboard Layout and Callbacks\n",
    "\n",
    "Plotly Dash Dashboards are composed of two parts. \n",
    "\n",
    "- The first part is the \"layout\" of the app and it describes what the application looks like. Dash apps are built off of a set of classes from the dash_core_components and dash_html_components libraries. These components are then arranged in a hierarchy of Python objects representing their layout.\n",
    "\n",
    "- The second part describes the interactivity of the application. This second part is written as a Python callback function. Callbacks are registered to update the application's layout when the values of the component's properties change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8377d7ad-7b75-4fea-8a45-bb7d94eee136",
   "metadata": {
    "id": "8377d7ad-7b75-4fea-8a45-bb7d94eee136"
   },
   "outputs": [],
   "source": [
    "# function to load dataset based on the selected dropdown value\n",
    "def read_dataset(state_name, gpu_enabled, currently_loaded_state, dtype_changed=False):\n",
    "    global c_df\n",
    "    if state_name != currently_loaded_state or dtype_changed:\n",
    "        if state_name == \"USA\":\n",
    "            data_path = f\"{DATA_PATH}/total_population_dataset.parquet\"\n",
    "        else:\n",
    "            data_path = f\"{DATA_PATH_STATE}/{state_name}.parquet\"\n",
    "        c_df = load_dataset(data_path, \"cudf\" if gpu_enabled else \"pandas\")\n",
    "    return c_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5f495c2",
   "metadata": {},
   "source": [
    "### Layout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8703572-b3e4-43cf-97d4-d2fc81f90ae9",
   "metadata": {
    "id": "b8703572-b3e4-43cf-97d4-d2fc81f90ae9"
   },
   "outputs": [],
   "source": [
    "# Describe the dashboards layout\n",
    "app.layout = html.Div(\n",
    "    children=[\n",
    "        html.Div(\n",
    "            children=[\n",
    "                html.H1(\n",
    "                    children=[\n",
    "                        \"Census 2020 Net Migration Visualization\",\n",
    "                        html.A(\n",
    "                            html.Img(\n",
    "                                src=\"assets/rapids-logo.png\",\n",
    "                                style={\n",
    "                                    \"float\": \"right\",\n",
    "                                    \"height\": \"45px\",\n",
    "                                    \"marginRight\": \"1%\",\n",
    "                                    \"marginTop\": \"-7px\",\n",
    "                                },\n",
    "                            ),\n",
    "                            href=\"https://rapids.ai/\",\n",
    "                        ),\n",
    "                        html.A(\n",
    "                            html.Img(\n",
    "                                src=\"assets/dash-logo.png\",\n",
    "                                style={\"float\": \"right\", \"height\": \"30px\"},\n",
    "                            ),\n",
    "                            href=\"https://dash.plot.ly/\",\n",
    "                        ),\n",
    "                    ],\n",
    "                    style={\"textAlign\": \"left\"},\n",
    "                ),\n",
    "            ]\n",
    "        ),\n",
    "        html.Div(\n",
    "            children=[\n",
    "                html.Div(\n",
    "                    children=[\n",
    "                        html.Div(\n",
    "                            children=[\n",
    "                                html.H4(\n",
    "                                    [\n",
    "                                        \"Population Count and Query Time\",\n",
    "                                    ],\n",
    "                                    className=\"container_title\",\n",
    "                                ),\n",
    "                                dcc.Loading(\n",
    "                                    dcc.Graph(\n",
    "                                        id=\"indicator-graph\",\n",
    "                                        figure=blank_fig(row_heights[3]),\n",
    "                                        config={\"displayModeBar\": False},\n",
    "                                    ),\n",
    "                                    color=\"#b0bec5\",\n",
    "                                    style={\"height\": f\"{row_heights[3]}px\"},\n",
    "                                ),\n",
    "                            ],\n",
    "                            style={\"height\": f\"{row_heights[0]}px\"},\n",
    "                            className=\"five columns pretty_container\",\n",
    "                            id=\"indicator-div\",\n",
    "                        ),\n",
    "                        html.Div(\n",
    "                            children=[\n",
    "                                html.Div(\n",
    "                                    children=[\n",
    "                                        html.Button(\n",
    "                                            \"Clear All Selections\",\n",
    "                                            id=\"clear-all\",\n",
    "                                            className=\"reset-button\",\n",
    "                                        ),\n",
    "                                    ]\n",
    "                                ),\n",
    "                                html.H4(\n",
    "                                    [\n",
    "                                        \"Options\",\n",
    "                                    ],\n",
    "                                    className=\"container_title\",\n",
    "                                ),\n",
    "                                html.Table(\n",
    "                                    [\n",
    "                                        html.Tr(\n",
    "                                            [\n",
    "                                                html.Td(\n",
    "                                                    html.Div(\"GPU Acceleration\"),\n",
    "                                                    className=\"config-label\",\n",
    "                                                ),\n",
    "                                                html.Td(\n",
    "                                                    html.Div(\n",
    "                                                        [\n",
    "                                                            daq.DarkThemeProvider(\n",
    "                                                                daq.BooleanSwitch(\n",
    "                                                                    on=True,\n",
    "                                                                    color=\"#00cc96\",\n",
    "                                                                    id=\"gpu-toggle\",\n",
    "                                                                )\n",
    "                                                            ),\n",
    "                                                            dbc.Tooltip(\n",
    "                                                                \"Caution: Using CPU compute for more than 50 million points is not recommended.\",\n",
    "                                                                target=\"gpu-toggle\",\n",
    "                                                                placement=\"bottom\",\n",
    "                                                                autohide=True,\n",
    "                                                                style={\n",
    "                                                                    \"textAlign\": \"left\",\n",
    "                                                                    \"fontSize\": \"15px\",\n",
    "                                                                    \"color\": \"white\",\n",
    "                                                                    \"width\": \"350px\",\n",
    "                                                                    \"padding\": \"15px\",\n",
    "                                                                    \"borderRadius\": \"5px\",\n",
    "                                                                    \"backgroundColor\": \"#2a2a2e\",\n",
    "                                                                },\n",
    "                                                            ),\n",
    "                                                        ]\n",
    "                                                    )\n",
    "                                                ),\n",
    "                                                #######  State Selection Dropdown ######\n",
    "                                                html.Td(\n",
    "                                                    html.Div(\"Select State\"),\n",
    "                                                    style={\"fontSize\": \"20px\"},\n",
    "                                                ),\n",
    "                                                html.Td(\n",
    "                                                    dcc.Dropdown(\n",
    "                                                        id=\"state-dropdown\",\n",
    "                                                        options=[\n",
    "                                                            {\"label\": i, \"value\": i}\n",
    "                                                            for i in state_names\n",
    "                                                        ],\n",
    "                                                        value=\"CALIFORNIA\",\n",
    "                                                    ),\n",
    "                                                    style={\n",
    "                                                        \"width\": \"25%\",\n",
    "                                                        \"height\": \"15px\",\n",
    "                                                    },\n",
    "                                                ),\n",
    "                                                ###### VIEWS ARE HERE ###########\n",
    "                                                html.Td(\n",
    "                                                    html.Div(\"Data-Selection\"),\n",
    "                                                    style={\"fontSize\": \"20px\"},\n",
    "                                                ),  # className=\"config-label\"\n",
    "                                                html.Td(\n",
    "                                                    dcc.Dropdown(\n",
    "                                                        id=\"view-dropdown\",\n",
    "                                                        options=[\n",
    "                                                            {\n",
    "                                                                \"label\": \"Total Population\",\n",
    "                                                                \"value\": \"total\",\n",
    "                                                            },\n",
    "                                                            {\n",
    "                                                                \"label\": \"Migrating In\",\n",
    "                                                                \"value\": \"in\",\n",
    "                                                            },\n",
    "                                                            {\n",
    "                                                                \"label\": \"Stationary\",\n",
    "                                                                \"value\": \"stationary\",\n",
    "                                                            },\n",
    "                                                            {\n",
    "                                                                \"label\": \"Migrating Out\",\n",
    "                                                                \"value\": \"out\",\n",
    "                                                            },\n",
    "                                                            {\n",
    "                                                                \"label\": \"Net Migration\",\n",
    "                                                                \"value\": \"net\",\n",
    "                                                            },\n",
    "                                                            {\n",
    "                                                                \"label\": \"Population with Race\",\n",
    "                                                                \"value\": \"race\",\n",
    "                                                            },\n",
    "                                                        ],\n",
    "                                                        value=\"net\",\n",
    "                                                        searchable=False,\n",
    "                                                        clearable=False,\n",
    "                                                    ),\n",
    "                                                    style={\n",
    "                                                        \"width\": \"25%\",\n",
    "                                                        \"height\": \"15px\",\n",
    "                                                    },\n",
    "                                                ),\n",
    "                                            ]\n",
    "                                        ),\n",
    "                                    ],\n",
    "                                    style={\"width\": \"100%\", \"marginTop\": \"30px\"},\n",
    "                                ),\n",
    "                                # Hidden div inside the app that stores the intermediate value\n",
    "                                html.Div(\n",
    "                                    id=\"datapoints-state-value\",\n",
    "                                    style={\"display\": \"none\"},\n",
    "                                ),\n",
    "                            ],\n",
    "                            style={\"height\": f\"{row_heights[0]}px\"},\n",
    "                            className=\"seven columns pretty_container\",\n",
    "                            id=\"config-div\",\n",
    "                        ),\n",
    "                    ]\n",
    "                ),\n",
    "                ##################### Map starts  ###################################\n",
    "                html.Div(\n",
    "                    children=[\n",
    "                        html.Button(\n",
    "                            \"Clear Selection\", id=\"reset-map\", className=\"reset-button\"\n",
    "                        ),\n",
    "                        html.H4(\n",
    "                            [\n",
    "                                \"Population Distribution of Individuals\",\n",
    "                            ],\n",
    "                            className=\"container_title\",\n",
    "                        ),\n",
    "                        dcc.Graph(\n",
    "                            id=\"map-graph\",\n",
    "                            config={\"displayModeBar\": False},\n",
    "                            figure=blank_fig(row_heights[1]),\n",
    "                        ),\n",
    "                        # Hidden div inside the app that stores the intermediate value\n",
    "                        html.Div(\n",
    "                            id=\"intermediate-state-value\", style={\"display\": \"none\"}\n",
    "                        ),\n",
    "                    ],\n",
    "                    className=\"twelve columns pretty_container\",\n",
    "                    id=\"map-div\",\n",
    "                    style={\"height\": \"50%\"},\n",
    "                ),\n",
    "                ################# Bars start #########################\n",
    "                # Race start\n",
    "                html.Div(\n",
    "                    children=[\n",
    "                        html.Div(\n",
    "                            children=[\n",
    "                                html.Button(\n",
    "                                    \"Clear Selection\",\n",
    "                                    id=\"clear-race\",\n",
    "                                    className=\"reset-button\",\n",
    "                                ),\n",
    "                                html.H4(\n",
    "                                    [\n",
    "                                        \"Race Distribution\",\n",
    "                                    ],\n",
    "                                    className=\"container_title\",\n",
    "                                ),\n",
    "                                dcc.Graph(\n",
    "                                    id=\"race-histogram\",\n",
    "                                    config={\"displayModeBar\": False},\n",
    "                                    figure=blank_fig(row_heights[2]),\n",
    "                                ),\n",
    "                            ],\n",
    "                            className=\"one-third column pretty_container\",\n",
    "                            id=\"race-div\",\n",
    "                        ),  # County top starts\n",
    "                        html.Div(\n",
    "                            children=[\n",
    "                                html.Button(\n",
    "                                    \"Clear Selection\",\n",
    "                                    id=\"clear-county-top\",\n",
    "                                    className=\"reset-button\",\n",
    "                                ),\n",
    "                                html.H4(\n",
    "                                    [\n",
    "                                        \"County-wise Top 15\",\n",
    "                                    ],\n",
    "                                    className=\"container_title\",\n",
    "                                ),\n",
    "                                dcc.Graph(\n",
    "                                    id=\"county-histogram-top\",\n",
    "                                    config={\"displayModeBar\": False},\n",
    "                                    figure=blank_fig(row_heights[2]),\n",
    "                                    animate=False,\n",
    "                                ),\n",
    "                            ],\n",
    "                            className=\" one-third column pretty_container\",\n",
    "                            id=\"county-div-top\",\n",
    "                        ),\n",
    "                        # County bottom starts\n",
    "                        html.Div(\n",
    "                            children=[\n",
    "                                html.Button(\n",
    "                                    \"Clear Selection\",\n",
    "                                    id=\"clear-county-bottom\",\n",
    "                                    className=\"reset-button\",\n",
    "                                ),\n",
    "                                html.H4(\n",
    "                                    [\n",
    "                                        \"County-wise Bottom 15\",\n",
    "                                    ],\n",
    "                                    className=\"container_title\",\n",
    "                                ),\n",
    "                                dcc.Graph(\n",
    "                                    id=\"county-histogram-bottom\",\n",
    "                                    config={\"displayModeBar\": False},\n",
    "                                    figure=blank_fig(row_heights[2]),\n",
    "                                    animate=False,\n",
    "                                ),\n",
    "                            ],\n",
    "                            className=\"one-third column pretty_container\",\n",
    "                        ),\n",
    "                    ],\n",
    "                    className=\"twelve columns\",\n",
    "                )\n",
    "                ############## End of  Bars #####################\n",
    "            ]\n",
    "        ),\n",
    "        html.Div(\n",
    "            [\n",
    "                html.H4(\"Acknowledgements and Data Sources\", style={\"marginTop\": \"0\"}),\n",
    "                dcc.Markdown(\n",
    "                    \"\"\"\\\n",
    "- 2020 Population Census and 2010 Population Census to compute Migration Dataset, used with permission from IPUMS NHGIS, University of Minnesota, [www.nhgis.org](https://www.nhgis.org/) ( not for redistribution ).\n",
    "- Base map layer provided by [Mapbox](https://www.mapbox.com/).\n",
    "- Dashboard developed with [Plotly Dash](https://plotly.com/dash/).\n",
    "- Geospatial point rendering developed with [Datashader](https://datashader.org/).\n",
    "- GPU toggle accelerated with [RAPIDS cudf and dask_cudf](https://rapids.ai/) and [cupy](https://cupy.chainer.org/), CPU toggle with [pandas](https://pandas.pydata.org/).\n",
    "- For source code and data workflow, visit our [GitHub](https://github.com/rapidsai/plotly-dash-rapids-census-demo/tree/master).\n",
    "\"\"\"\n",
    "                ),\n",
    "            ],\n",
    "            style={\n",
    "                \"width\": \"98%\",\n",
    "                \"marginRight\": \"0\",\n",
    "                \"padding\": \"10px\",\n",
    "            },\n",
    "            className=\"twelve columns pretty_container\",\n",
    "        ),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4407d4f",
   "metadata": {},
   "source": [
    "### Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03dc058b",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Clear/reset button callbacks\n",
    "@app.callback(\n",
    "    Output(\"map-graph\", \"selectedData\"),\n",
    "    [Input(\"reset-map\", \"n_clicks\"), Input(\"clear-all\", \"n_clicks\")],\n",
    ")\n",
    "def clear_map(*args):\n",
    "    return None\n",
    "\n",
    "\n",
    "@app.callback(\n",
    "    Output(\"race-histogram\", \"selectedData\"),\n",
    "    [Input(\"clear-race\", \"n_clicks\"), Input(\"clear-all\", \"n_clicks\")],\n",
    ")\n",
    "def clear_race_hist_selections(*args):\n",
    "    return None\n",
    "\n",
    "\n",
    "@app.callback(\n",
    "    Output(\"county-histogram-top\", \"selectedData\"),\n",
    "    [Input(\"clear-county-top\", \"n_clicks\"), Input(\"clear-all\", \"n_clicks\")],\n",
    ")\n",
    "def clear_county_hist_top_selections(*args):\n",
    "    return None\n",
    "\n",
    "\n",
    "@app.callback(\n",
    "    Output(\"county-histogram-bottom\", \"selectedData\"),\n",
    "    [Input(\"clear-county-bottom\", \"n_clicks\"), Input(\"clear-all\", \"n_clicks\")],\n",
    ")\n",
    "def clear_county_hist_bottom_selections(*args):\n",
    "    return None\n",
    "\n",
    "\n",
    "@app.callback(\n",
    "    [\n",
    "        Output(\"indicator-graph\", \"figure\"),\n",
    "        Output(\"map-graph\", \"figure\"),\n",
    "        Output(\"map-graph\", \"config\"),\n",
    "        Output(\"county-histogram-top\", \"figure\"),\n",
    "        Output(\"county-histogram-top\", \"config\"),\n",
    "        Output(\"county-histogram-bottom\", \"figure\"),\n",
    "        Output(\"county-histogram-bottom\", \"config\"),\n",
    "        Output(\"race-histogram\", \"figure\"),\n",
    "        Output(\"race-histogram\", \"config\"),\n",
    "        Output(\"intermediate-state-value\", \"children\"),\n",
    "    ],\n",
    "    [\n",
    "        Input(\"map-graph\", \"relayoutData\"),\n",
    "        Input(\"map-graph\", \"selectedData\"),\n",
    "        Input(\"race-histogram\", \"selectedData\"),\n",
    "        Input(\"county-histogram-top\", \"selectedData\"),\n",
    "        Input(\"county-histogram-bottom\", \"selectedData\"),\n",
    "        Input(\"view-dropdown\", \"value\"),\n",
    "        Input(\"state-dropdown\", \"value\"),\n",
    "        Input(\"gpu-toggle\", \"on\"),\n",
    "    ],\n",
    "    [\n",
    "        State(\"intermediate-state-value\", \"children\"),\n",
    "    ],\n",
    "    # prevent_initial_call=True\n",
    ")\n",
    "def update_plots(\n",
    "    relayout_data,\n",
    "    selected_map,\n",
    "    selected_race,\n",
    "    selected_county_top,\n",
    "    selected_county_bottom,\n",
    "    view_name,\n",
    "    state_name,\n",
    "    gpu_enabled,\n",
    "    coordinates_backup,\n",
    "):\n",
    "    global data_3857, data_center_3857, data_4326, data_center_4326\n",
    "    global currently_loaded_state, selected_race_backup, selected_county_top_backup, selected_county_bt_backup\n",
    "\n",
    "    # condition to avoid reloading on tool update\n",
    "    if (\n",
    "        ctx.triggered_id == \"map-graph\"\n",
    "        and relayout_data\n",
    "        and list(relayout_data.keys()) == [\"dragmode\"]\n",
    "    ):\n",
    "        raise PreventUpdate\n",
    "\n",
    "    # condition to avoid a bug in plotly where selectedData is reset following a box-select\n",
    "    if not (selected_race is not None and len(selected_race[\"points\"]) == 0):\n",
    "        selected_race_backup = selected_race\n",
    "    elif ctx.triggered_id == \"race-histogram\":\n",
    "        raise PreventUpdate\n",
    "\n",
    "    # condition to avoid a bug in plotly where selectedData is reset following a box-select\n",
    "    if not (\n",
    "        selected_county_top is not None and len(selected_county_top[\"points\"]) == 0\n",
    "    ):\n",
    "        selected_county_top_backup = selected_county_top\n",
    "    elif ctx.triggered_id == \"county-histogram-top\":\n",
    "        raise PreventUpdate\n",
    "\n",
    "    # condition to avoid a bug in plotly where selectedData is reset following a box-select\n",
    "    if not (\n",
    "        selected_county_bottom is not None\n",
    "        and len(selected_county_bottom[\"points\"]) == 0\n",
    "    ):\n",
    "        selected_county_bt_backup = selected_county_bottom\n",
    "    elif ctx.triggered_id == \"county-histogram-bottom\":\n",
    "        raise PreventUpdate\n",
    "\n",
    "    df = read_dataset(state_name, gpu_enabled, currently_loaded_state, dtype_changed=ctx.triggered_id == \"gpu-toggle\")\n",
    "\n",
    "    t0 = time.time()\n",
    "\n",
    "    if coordinates_backup is not None:\n",
    "        coordinates_4326_backup, position_backup = coordinates_backup\n",
    "    else:\n",
    "        coordinates_4326_backup, position_backup = None, None\n",
    "\n",
    "    colorscale_name = \"Viridis\"\n",
    "\n",
    "    if data_3857 == [] or state_name != currently_loaded_state:\n",
    "        (\n",
    "            data_3857,\n",
    "            data_center_3857,\n",
    "            data_4326,\n",
    "            data_center_4326,\n",
    "        ) = set_projection_bounds(df)\n",
    "\n",
    "    (\n",
    "        datashader_plot,\n",
    "        race_histogram,\n",
    "        county_top_histogram,\n",
    "        county_bottom_histogram,\n",
    "        n_selected_indicator,\n",
    "        coordinates_4326_backup,\n",
    "        position_backup,\n",
    "    ) = build_updated_figures(\n",
    "        df,\n",
    "        relayout_data,\n",
    "        selected_map,\n",
    "        selected_race_backup,\n",
    "        selected_county_top_backup,\n",
    "        selected_county_bt_backup,\n",
    "        colorscale_name,\n",
    "        data_3857,\n",
    "        data_center_3857,\n",
    "        data_4326,\n",
    "        data_center_4326,\n",
    "        coordinates_4326_backup,\n",
    "        position_backup,\n",
    "        view_name,\n",
    "    )\n",
    "\n",
    "    barchart_config = {\n",
    "        \"displayModeBar\": True,\n",
    "        \"modeBarButtonsToRemove\": [\n",
    "            \"zoom2d\",\n",
    "            \"pan2d\",\n",
    "            \"select2d\",\n",
    "            \"lasso2d\",\n",
    "            \"zoomIn2d\",\n",
    "            \"zoomOut2d\",\n",
    "            \"resetScale2d\",\n",
    "            \"hoverClosestCartesian\",\n",
    "            \"hoverCompareCartesian\",\n",
    "            \"toggleSpikelines\",\n",
    "        ],\n",
    "    }\n",
    "    compute_time = time.time() - t0\n",
    "    print(f\"Query time: {compute_time}\")\n",
    "    n_selected_indicator[\"data\"].append(\n",
    "        {\n",
    "            \"title\": {\"text\": \"Query Time\"},\n",
    "            \"type\": \"indicator\",\n",
    "            \"value\": round(compute_time, 4),\n",
    "            \"domain\": {\"x\": [0.6, 0.85], \"y\": [0, 0.5]},\n",
    "            \"number\": {\n",
    "                \"font\": {\n",
    "                    \"color\": text_color,\n",
    "                    \"size\": \"50px\",\n",
    "                },\n",
    "                \"suffix\": \" seconds\",\n",
    "            },\n",
    "        }\n",
    "    )\n",
    "\n",
    "    datashader_plot[\"layout\"][\"dragmode\"] = (\n",
    "        relayout_data[\"dragmode\"]\n",
    "        if (relayout_data and \"dragmode\" in relayout_data)\n",
    "        else dragmode_backup\n",
    "    )\n",
    "    # update currently loaded state\n",
    "    currently_loaded_state = state_name\n",
    "\n",
    "    return (\n",
    "        n_selected_indicator,\n",
    "        datashader_plot,\n",
    "        {\n",
    "            \"displayModeBar\": True,\n",
    "            \"modeBarButtonsToRemove\": [\n",
    "                \"lasso2d\",\n",
    "                \"zoomInMapbox\",\n",
    "                \"zoomOutMapbox\",\n",
    "                \"toggleHover\",\n",
    "            ],\n",
    "        },\n",
    "        county_top_histogram,\n",
    "        barchart_config,\n",
    "        county_bottom_histogram,\n",
    "        barchart_config,\n",
    "        race_histogram,\n",
    "        barchart_config,\n",
    "        (coordinates_4326_backup, position_backup),\n",
    "    )\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c4e7c19-f12e-4811-b897-8df7a2633ca3",
   "metadata": {},
   "source": [
    "> **Note: If you encounter unknown errors when running the app, it is recommended to select \"Restart runtime and run all\" from the \"Runtime\" menu option on top. Additionally, you can comment out the installation commands at the top of the file. Colab allows for a persistent environment even when you restart the runtime.**\n",
    "\n",
    "> **Under ~10M rows the GPU speed up over CPU will be negligible but increase dramatically**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83b18e1f",
   "metadata": {},
   "source": [
    "### Run the Plotly Dash app"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "r_UVM7AvC_4v",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 708
    },
    "id": "r_UVM7AvC_4v",
    "outputId": "fa79ff95-4ab6-489b-95c1-32707820a47c"
   },
   "outputs": [],
   "source": [
    "app.run_server(debug=False)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
